require(ggplot2)
require(reshape2)
require(viridis)

delta <- rho <- seq(0.05, 1, length.out = length.out )

sim_res <- list()
for (i in 1:length(delta)) {
  fn <- paste("./change/corr07/pred/res/sim_pred_change_mse_corr07_", i, ".rds", sep = '')
  sim_res[[i]] <-  readRDS(fn)
}

## data prep for plot
mse_lasso1 <- do.call("cbind", lapply(sim_res, function(x) x[,1]))
mse_lasso2 <- do.call("cbind", lapply(sim_res, function(x) x[,2]))
mse_bridge <- do.call("cbind", lapply(sim_res, function(x) x[,3]))

mse_lasso1_dat <- data.frame(mse_lasso1)
mse_lasso2_dat <- data.frame(mse_lasso2)
mse_bridge_dat <- data.frame(mse_bridge)

colnames(mse_lasso1_dat) <- colnames(mse_lasso2_dat) <- colnames(mse_bridge_dat)  <- paste("delta", delta, sep = '')
mse_lasso1_dat$rho  <- rho
mse_lasso2_dat$rho  <- rho
mse_bridge_dat$rho  <- rho

mse_lasso1_dat$method <- "HMM Lasso"
mse_lasso2_dat$method <- "Oracle Lasso"
mse_bridge_dat$method <- "HMBB"

mse_full <- rbind(mse_lasso1_dat, mse_lasso2_dat, mse_bridge_dat)

mse_full <- melt(mse_full, id.vars = c("rho", "method"), variable.name = "delta", value.name = 'L2')
mse_full$delta <- delta[as.numeric(mse_full$delta)]



mse_full$L21 <- round(mse_full$L2, 3)
# fig_save <- ggplot(data = subset(mse_full, method != "Ridge"), aes(x = delta, y = rho)) + geom_tile(aes(fill = L2)) +
#     scale_fill_distiller(palette = "YlGnBu", direction = 1) +
#     # scale_fill_viridis(option = "magma", direction = -1) +
#     labs(x = expression(paste(alpha, " = ", N / p)), y  = expression(paste(rho, " = ", k / N)), fill = "") +
#     scale_x_continuous(expand = c(0,0)) +
#     scale_y_continuous(expand = c(0,0)) +
#     facet_wrap(~method) +
#     theme_bw() +
#     theme(panel.background = element_blank(),
#             panel.grid = element_blank(),
#             panel.border = element_rect(color = 'gray30', size = 0.7),
#             strip.background = element_blank(),
#             strip.text = element_text(face = 'bold', size = 12),
#             axis.text.x = element_text(size = 11, color = 'gray30'),
#             axis.text.y = element_text(size = 11, color = 'gray30'),
#             axis.title.x = element_text(size = 12.5),
#             axis.title.y = element_text(size = 12.5))

# fig_save <- ggplot(mse_full, aes(x = delta, y = rho)) + geom_tile(aes(fill = L21)) +
#     scale_fill_distiller(palette = "YlGnBu", direction = 1) +
#     # scale_fill_viridis(option = "magma", direction = -1) +
#     labs(x = expression(paste(delta, " = ", N / p)), y  = expression(paste(rho, " = ", k / N)), fill = "") +
#     scale_x_continuous(expand = c(0,0)) +
#     scale_y_continuous(expand = c(0,0)) +
#     facet_wrap(~method) +
#     theme_bw() +
#     theme(panel.background = element_blank(),
#             panel.grid = element_blank(),
#             panel.border = element_rect(color = 'gray30', size = 0.7),
#             strip.background = element_blank(),
#             strip.text = element_text(face = 'bold', size = 12),
#             axis.text.x = element_text(size = 11, color = 'gray30'),
#             axis.text.y = element_text(size = 11, color = 'gray30'),
#             axis.title.x = element_text(size = 12.5),
#             axis.title.y = element_text(size = 12.5))
#
# ggsave(fig_save, file = 'L2_pred_error_corr0.pdf', height = 6.5, width = 7)

fig_save <- ggplot(mse_full, aes(x = delta, y = rho)) + geom_tile(aes(fill = L21)) +
    # scale_fill_distiller(palette = "YlGnBu", direction = 1) +
    scale_fill_distiller(palette = "RdYlBu") +
    # scale_fill_viridis(option = "magma", direction = -1) +
    labs(x = expression(paste(delta, " = ", n / p)), y  = expression(paste(rho, " = ", k / n)), fill = expression(paste("x", 10^2))) +
    scale_x_continuous(expand = c(0,0)) +
    scale_y_continuous(expand = c(0,0)) +
    facet_wrap(~method, ncol = 4) +
    ggtitle("Panel A (Prediction Loss)") +
    theme_bw() +
    theme(panel.background = element_blank(),
            panel.grid = element_blank(),
            panel.border = element_rect(color = 'gray30', size = 0.7),
            plot.title = element_text(face = 'bold'),
            strip.background = element_blank(),
            strip.text = element_text(face = 'bold', size = 12),
            axis.text.x = element_text(size = 11, color = 'gray30'),
            axis.text.y = element_text(size = 11, color = 'gray30'),
            axis.title.x = element_text(size = 12.5),
            axis.title.y = element_text(size = 12.5))

ggsave(fig_save, file = 'SI/Figure7_Panel_A.pdf', height = 2.7, width = 6)
